package com.snow.util;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树构建工具
 *
 * @author by 13119
 */
public class TreeUtil {
    private TreeUtil() {

    }

    /**
     * 树构建工具
     * @param surplus 需要构建的列表
     * @param rootKey 根节点值
     * @param getId 获取跟节点值的方法
     * @param getParentId 获取节点父【key】值的方法
     * @param setChild 设置字子点的方法
     * @return 构建成功的list
     * @param <E> 构建类型
     * @param <K> 节点分组主键【key】
     */
    public static <E, K> List<E> buildTree(List<E> surplus, K rootKey, Function<E, K> getId, Function<E, K> getParentId, BiConsumer<E, List<E>> setChild) {
        return buildTree(surplus,rootKey,getId,getParentId,setChild,null);
    }
    public static <E, K> List<E> buildTree(List<E> surplus, K rootKey, Function<E, K> getId, Function<E, K> getParentId, BiConsumer<E, List<E>> setChild,Comparator<? super E> sortComparator) {
        if (surplus == null || surplus.isEmpty()) {
            return new ArrayList<>();
        }
        // 过滤根节点
        Map<String, List<E>> groupResult = surplus.stream().collect(Collectors.groupingBy(item -> getParentId.apply(item).equals(rootKey) ? "root" : "surplus"));
        //获取根节点
        List<E> roots = groupResult.get("root");
        if (sortComparator != null) {
            roots.sort(sortComparator);
        }
        //获取剩余节点
        surplus = groupResult.get("surplus");
        //未找到符合条件的根节点直接返回
        if (roots == null) {
            return new ArrayList<>();
        }
        //其他节点为空，直接返回
        if (surplus == null) {
            return roots;
        }
        for (E node : roots) {
            List<E> child = buildTree(surplus, getId.apply(node), getId, getParentId,setChild);
            if (child.isEmpty()) {
                return roots;
            }
            if (sortComparator != null) {
                child.sort(sortComparator);
            }
            setChild.accept(node,child);
        }
        return roots;
    }
}
